#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>

#define CHANNELS 3

__global__ void directConvolution(float *output, float *input, float *kernel, 
                                 int input_size, int kernel_size, int stride, 
                                 int output_size, int padding) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    int c = blockIdx.z * blockDim.z + threadIdx.z;

    if (x < output_size && y < output_size && c < CHANNELS) {
        float sum = 0.0f;
        
        for (int kx = 0; kx < kernel_size; ++kx) {
            for (int ky = 0; ky < kernel_size; ++ky) {
                int input_x = x * stride + kx - padding;
                int input_y = y * stride + ky - padding;
                
                if (input_x >= 0 && input_x < input_size && 
                    input_y >= 0 && input_y < input_size) {
                    for (int kc = 0; kc < CHANNELS; ++kc) {
                        float input_val = input[(input_y * input_size + input_x) * CHANNELS + kc];
                        float kernel_val = kernel[(ky * kernel_size + kx) * CHANNELS * CHANNELS + kc * CHANNELS + c];
                        sum += input_val * kernel_val;
                    }
                }
            }
        }
        
        output[(y * output_size + x) * CHANNELS + c] = sum;
    }
}

int main(int argc, char **argv) {
    if (argc != 4) {
        printf("Usage: %s <input_size> <kernel_size> <stride>\n", argv[0]);
        return 1;
    }

    int input_size = atoi(argv[1]);
    int kernel_size = atoi(argv[2]);
    int stride = atoi(argv[3]);
    int padding = (kernel_size - 1) / 2;
    int output_size = (input_size + 2 * padding - kernel_size) / stride + 1;

    // Allocate and initialize host memory
    size_t input_bytes = input_size * input_size * CHANNELS * sizeof(float);
    size_t kernel_bytes = kernel_size * kernel_size * CHANNELS * CHANNELS * sizeof(float);
    size_t output_bytes = output_size * output_size * CHANNELS * sizeof(float);

    float *h_input = (float *)malloc(input_bytes);
    float *h_kernel = (float *)malloc(kernel_bytes);
    float *h_output = (float *)malloc(output_bytes);

    // Initialize with random values
    for (int i = 0; i < input_size * input_size * CHANNELS; i++)
        h_input[i] = (float)rand() / RAND_MAX;
    for (int i = 0; i < kernel_size * kernel_size * CHANNELS * CHANNELS; i++)
        h_kernel[i] = (float)rand() / RAND_MAX;

    // Allocate device memory
    float *d_input, *d_kernel, *d_output;
    cudaMalloc(&d_input, input_bytes);
    cudaMalloc(&d_kernel, kernel_bytes);
    cudaMalloc(&d_output, output_bytes);

    // Copy data to device
    cudaMemcpy(d_input, h_input, input_bytes, cudaMemcpyHostToDevice);
    cudaMemcpy(d_kernel, h_kernel, kernel_bytes, cudaMemcpyHostToDevice);

    // Launch kernel
    dim3 threads(16, 16, 1);
    dim3 blocks((output_size + threads.x - 1) / threads.x,
                (output_size + threads.y - 1) / threads.y,
                CHANNELS);

    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    cudaEventRecord(start);

    directConvolution<<<blocks, threads>>>(d_output, d_input, d_kernel, 
                                         input_size, kernel_size, stride, 
                                         output_size, padding);

    cudaEventRecord(stop);
    cudaEventSynchronize(stop);

    float milliseconds = 0;
    cudaEventElapsedTime(&milliseconds, start, stop);

    // Copy result back to host
    cudaMemcpy(h_output, d_output, output_bytes, cudaMemcpyDeviceToHost);

    printf("Direct Convolution Time: %.2f ms\n", milliseconds);

    // Cleanup
    free(h_input);
    free(h_kernel);
    free(h_output);
    cudaFree(d_input);
    cudaFree(d_kernel);
    cudaFree(d_output);

    return 0;
}